#include "VulkanTexture2Ds.h"
#include "VulkanTexture2D.h"
#include "Render/VulkanContext.h"

#include "Math/MathBase.h"
#include "Render/Texture/VulkanTextureFormat.h"
#include "Render/Texture/Texture2DPixels.h"
#include "Render/PreRender/PreRenderTexture2D.h"



void CVulkanTexture2Ds::UpdateTextures(const TVector<std::shared_ptr<CPreRenderTexture2D>> &AllTextureData)
{
    auto OldTexture2Ds=Texture2Ds;
    
    Texture2Ds.Clear();
    TVector< TThreadSafeSharedPtr<CTexture2DPixels>> NeedUpdateTextures;
    for (const auto& PreRenderTexture : AllTextureData)
    {
        auto TextureData=PreRenderTexture->Pixels;
        if(OldTexture2Ds.Contains(TextureData))
        {
            Texture2Ds.Add(TextureData, OldTexture2Ds[TextureData]);
            if(TextureData->IsDirty())
            {
                NeedUpdateTextures.Add(TextureData);
            }

            continue;
        }
        else
        {
            Texture2Ds.Add(TextureData, std::make_shared<CVulkanTexture2D>());
            NeedUpdateTextures.Add(TextureData);
        }

    }

    OldTexture2Ds.Clear();

    //更新纹理数据
    TVector<std::shared_ptr<CVulkanTexture2D>> NeedConvertLayoutTextures;
    for (const auto& TextureData : NeedUpdateTextures)
    {
        auto VulkanTexture=Texture2Ds[TextureData];
        std::lock_guard<std::mutex> LockGuard(TextureData->Mutex);
        TextureData->bDirty=false;

        vk::Format Format = ToVulkanTextureFormat(TextureData->Format);
        assert(Format != vk::Format::eUndefined);

        VulkanTexture->Image=CVulkanImage(); //重新创建一个图像，原有的图像会被回收
        VulkanTexture->Image.SetTexture2DImageParams(TextureData->Width, TextureData->Height, Format);
        VulkanTexture->Image.WriteData(TextureData->Data.data(), TextureData->Data.Num(), TextureData->Pitch);

        vk::ImageViewCreateInfo ImageViewCreateInfo=vk::ImageViewCreateInfo()
            .setImage(VulkanTexture->Image.GetImage())
            .setViewType(vk::ImageViewType::e2D)
            .setFormat(Format)
            .setSubresourceRange(vk::ImageSubresourceRange()
                .setAspectMask(vk::ImageAspectFlagBits::eColor)
                .setBaseMipLevel(0)
                .setLevelCount(1)
                .setBaseArrayLayer(0)
                .setLayerCount(1));

        //销毁旧的ImageView
        CVulkanContext::Get()->GetRecycleObjects()->ImageViews.Add(VulkanTexture->ImageView);
        VulkanTexture->ImageView=CVulkanContext::Get()->GetDevice().createImageView(ImageViewCreateInfo);
        VulkanTexture->Width=TextureData->Width;
        VulkanTexture->Height=TextureData->Height;

        NeedConvertLayoutTextures.Add(VulkanTexture);

    }

    if(!NeedConvertLayoutTextures.Empty())
    {

        //转换布局
        vk::CommandBuffer ConvertImageLayoutCommandBuffer=CVulkanContext::Get()->GetConvertImageLayoutCommandBuffer();

        vk::CommandBufferBeginInfo CommandBufferBeginInfo=vk::CommandBufferBeginInfo()
            .setFlags(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);

        ConvertImageLayoutCommandBuffer.begin(CommandBufferBeginInfo);

        for(auto VulkanTexture: NeedConvertLayoutTextures)
        {
            vk::ImageMemoryBarrier ImageMemoryBarrier=vk::ImageMemoryBarrier()
            .setSrcAccessMask(vk::AccessFlagBits::eHostWrite)
            .setDstAccessMask(vk::AccessFlagBits::eShaderRead)
            .setOldLayout(vk::ImageLayout::ePreinitialized)
            .setNewLayout(vk::ImageLayout::eShaderReadOnlyOptimal)
            .setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
            .setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
            .setImage(VulkanTexture->Image.GetImage())
            .setSubresourceRange(vk::ImageSubresourceRange()
                .setAspectMask(vk::ImageAspectFlagBits::eColor)
                .setBaseMipLevel(0)
                .setLevelCount(1)
                .setBaseArrayLayer(0)
                .setLayerCount(1));

            ConvertImageLayoutCommandBuffer.pipelineBarrier(
                vk::PipelineStageFlagBits::eHost
                , vk::PipelineStageFlagBits::eFragmentShader
                , vk::DependencyFlags()
                , nullptr
                , nullptr
                , ImageMemoryBarrier);

        }
      
        ConvertImageLayoutCommandBuffer.end();

        vk::SubmitInfo SubmitInfo=vk::SubmitInfo()
            .setCommandBufferCount(1)
            .setPCommandBuffers(&ConvertImageLayoutCommandBuffer);

        CVulkanContext::Get()->GetGraphicQueue().submit(SubmitInfo
            , CVulkanContext::Get()->GetConvertImageLayoutFence());

        //wait for convert image layout
        CVulkanContext::Get()->GetDevice().waitForFences(
            {CVulkanContext::Get()->GetConvertImageLayoutFence()}
            , VK_TRUE
            , std::numeric_limits<uint64_t>::max());

        CVulkanContext::Get()->GetDevice().resetFences(
            {CVulkanContext::Get()->GetConvertImageLayoutFence()});

    }
        
    
    
}

std::shared_ptr<CVulkanTexture2D> CVulkanTexture2Ds::FindVulkanObj(TThreadSafeSharedPtr<CTexture2DPixels> InWriteInTexture)
{
    return Texture2Ds[InWriteInTexture];
}
